package com.mwl.gateway.filter;

import com.mwl.gateway.constant.GlobalVariable;
import com.mwl.gateway.utils.MdcUtils;
import jakarta.annotation.Nonnull;
import jakarta.validation.constraints.NotNull;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.slf4j.MDC;
import org.springframework.core.Ordered;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.server.WebFilter;
import org.springframework.web.server.WebFilterChain;
import reactor.core.publisher.Mono;
import reactor.util.context.Context;

import java.util.Map;
import java.util.UUID;

/**
 * description:
 *
 * @author maweilong
 **/
public class TraceInfoFilter implements WebFilter, Ordered {
    private static final Logger log = LoggerFactory.getLogger(TraceInfoFilter.class);

    @Override
    public Mono<Void> filter(@Nonnull ServerWebExchange exchange,
                             @Nonnull WebFilterChain chain) {
        long nanoTime = System.nanoTime();
        String traceId = UUID.randomUUID().toString().replace("-", "");
        Map<String, String> mdcContextMap = Map.of(GlobalVariable.TRACE_ID_HEADER, traceId);
        MDC.setContextMap(mdcContextMap);
        exchange.getResponse().getHeaders().set(GlobalVariable.TRACE_ID_HEADER, traceId);
        return chain.filter(exchange)
                .then(
                        Mono.<Void>fromRunnable(() -> {
                                    MDC.setContextMap(mdcContextMap);
                                    long elapsed = System.nanoTime() - nanoTime;
                                    long micro = elapsed / 1000;
                                    double mills = micro / 1000D;
                                    MdcUtils.checkMDC(traceId, "MDC filter");
                                    log.info("{} Request completed in {} ms", traceId, mills);
                                })
                                .contextWrite(Context.of(GlobalVariable.MDC_KEY, mdcContextMap))
                )
                .contextWrite(Context.of(GlobalVariable.MDC_KEY, mdcContextMap));
    }

    @Override
    public int getOrder() {
        return Ordered.HIGHEST_PRECEDENCE + 1;
    }
}
